# 智能体节点
import os
import json
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from state import TranslationState
from typing import Dict, Any
from langchain_core.output_parsers import StrOutputParser
from prompts import preprocess_node_prompt, terminology_node_prompt, translate_node_prompt, review_node_prompt
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import StrOutputParser


load_dotenv()
# 模型 - 改用Ollama
llm = ChatOllama(
    model="qwen3:4b",  # 确保你已通过Ollama部署了这个模型
    base_url="http://localhost:11434"  # Ollama默认地址，根据实际部署情况修改
)

# 预处理节点
def preprocess_node(state: TranslationState) -> Dict[str, Any]:
    print("=== 文本预处理 ===")
    print(f"输入文本: {state['original_text']}")
    corrected = (preprocess_node_prompt | llm | StrOutputParser()).invoke({"original_text": state["original_text"]})
    print(f"预处理后文本: {corrected}")
    print("=== 文本预处理结束 ===\n")
    return {"preprocessed_text": corrected}


# 术语检查节点
def terminology_node(state: TranslationState) -> Dict[str, Any]:
    print("=== 开始术语检查 ===")
    print(f"输入文本: {state['preprocessed_text']}")
    terminology_json = (terminology_node_prompt | llm | StrOutputParser()).invoke({"preprocessed_text": state["preprocessed_text"]})
    print(f"识别的领域和术语: {terminology_json}")
    print("=== 术语检查结束 ===\n")
    return {"terminology_json": terminology_json}



# 翻译节点
def translate_node(state: TranslationState) -> Dict[str, Any]:
    print("=== 开始翻译 ===")
    print(f"预处理文本: {state['preprocessed_text']}")
    
    # 安全处理术语表JSON数据
    term_data = {"domain": "", "key_terms": []}  # 默认空数据
    terminology_json = state.get("terminology_json", "")
    
    if terminology_json:
        try:
            term_data = json.loads(terminology_json)
            # 验证必要的字段结构
            if "domain" not in term_data:
                term_data["domain"] = ""
            if "key_terms" not in term_data:
                term_data["key_terms"] = []
        except json.JSONDecodeError as e:
            print(f"术语表JSON解析错误: {str(e)}，将使用空术语表")
        except Exception as e:
            print(f"处理术语表时发生错误: {str(e)}，将使用空术语表")
    
    # 构建术语字符串
    terms_str = "\n".join(f"- {t['term']}: {t['explanation']}" for t in term_data["key_terms"])
    print(f"领域: {term_data['domain']}")
    print(f"术语表:\n{terms_str}")
    
    # 安全处理审阅意见
    review_comments = ""
    review_json = state.get("review_json")
    if review_json:
        try:
            review_data = json.loads(review_json)
            if "comments" in review_data and isinstance(review_data["comments"], list):
                review_comments = "\n".join(f"- {comment}" for comment in review_data["comments"])
        except json.JSONDecodeError as e:
            print(f"审阅意见JSON解析错误: {str(e)}，将忽略审阅意见")
        except Exception as e:
            print(f"处理审阅意见时发生错误: {str(e)}，将忽略审阅意见")
    
    if review_comments:
        print(f"审阅意见:\n{review_comments}")
    
    # 执行翻译
    trans = (translate_node_prompt | llm | StrOutputParser()).invoke({
        "domain": term_data['domain'], 
        "terms": terms_str, 
        "preprocessed_text": state["preprocessed_text"],
        "review_comments": review_comments
    })
    print(f"翻译结果: {trans}")
    print("=== 翻译结束 ===\n")
    return {"current_translation": trans}


def review_node(state: TranslationState) -> Dict[str, Any]:
    print("=== 开始审阅 ===")
    # 先验证关键状态值是否存在，避免KeyError
    original_text = state.get("original_text", "")
    current_translation = state.get("current_translation", "")
    print(f"原文: {original_text}")
    print(f"当前译文: {current_translation}")

    # 1. 调用LLM生成审阅结果，并捕获可能的调用异常
    review_json = ""
    try:
        # 调用LLM获取审阅输出（确保prompt已明确要求返回标准JSON）
        review_json = (review_node_prompt | llm | StrOutputParser()).invoke({
            "original_text": original_text, 
            "current_translation": current_translation
        })
        print(f"LLM原始输出: {review_json}")  # 调试用：查看LLM实际返回内容
    except Exception as e:
        print(f"LLM调用失败: {str(e)}")
        # 兜底：生成默认的JSON结构，避免后续解析崩溃
        review_json = '{"comments": ["LLM审阅调用失败，使用默认结果"], "continue": false}'

    # 2. 安全解析JSON，处理空值、格式错误
    review = {"comments": [], "continue": False}  # 默认兜底结构
    if review_json.strip():  # 先判断是否为空字符串
        try:
            review = json.loads(review_json)
            # 验证JSON必须包含的字段，缺失则补默认值
            if "comments" not in review or not isinstance(review["comments"], list):
                review["comments"] = ["审阅结果缺少有效意见列表，已自动补全"]
            if "continue" not in review or not isinstance(review["continue"], bool):
                review["continue"] = False  # 默认不继续迭代
        except json.JSONDecodeError as e:
            print(f"审阅JSON解析失败: {str(e)}（原始内容：{review_json}）")
            # 解析失败时兜底，避免流程中断
            review = {
                "comments": [f"审阅结果格式错误：{str(e)}，已使用默认结果"], 
                "continue": False
            }
        except Exception as e:
            print(f"处理审阅结果时意外错误: {str(e)}")
            review = {
                "comments": [f"审阅结果处理异常：{str(e)}，已使用默认结果"], 
                "continue": False
            }
    else:
        # 处理LLM返回空字符串的情况
        print("LLM返回空的审阅结果，使用默认兜底")
        review = {"comments": ["LLM未返回有效审阅结果，已自动补全"], "continue": False}

    # 3. 提取审阅意见（即使是兜底结果也能正常处理）
    review_comments = "\n".join(f"- {comment}" for comment in review["comments"])
    print(f"最终审阅意见:\n{review_comments}")
    print("=== 当前审阅结束 ===\n")

    # 4. 构建返回结果，确保字段完整性
    # 先获取当前迭代次数，避免state中无该字段导致KeyError
    current_iter = state.get("iteration", 0)
    max_iter = state.get("max_iterations", 3)  # 默认最大迭代3次，防止死循环

    result = {
        "review_json": json.dumps(review),  # 存储标准JSON（避免原始错误内容）
        "review_comments": review_comments,
        "should_continue": review["continue"],
        "iteration": current_iter + 1,
    }

    # 5. 判断是否设置最终译文（停止迭代条件）
    if not review["continue"] or (current_iter + 1) >= max_iter:
        result["final_translation"] = current_translation  # 使用当前译文作为最终结果

    return result
